package cn.apelx.shardingjdbc.config;

import lombok.RequiredArgsConstructor;
import org.apache.shardingsphere.driver.api.ShardingSphereDataSourceFactory;
import org.apache.shardingsphere.infra.config.RuleConfiguration;
import org.apache.shardingsphere.infra.config.mode.ModeConfiguration;
import org.apache.shardingsphere.spring.boot.ShardingSphereAutoConfiguration;
import org.apache.shardingsphere.spring.boot.datasource.DataSourceMapSetter;
import org.apache.shardingsphere.spring.boot.prop.SpringBootPropertiesConfiguration;
import org.apache.shardingsphere.spring.boot.schema.SchemaNameSetter;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.BeanFactory;
import org.springframework.beans.factory.BeanFactoryAware;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.beans.factory.support.DefaultListableBeanFactory;
import org.springframework.boot.autoconfigure.AutoConfigureAfter;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.EnvironmentAware;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.ComponentScan;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.DependsOn;
import org.springframework.core.env.Environment;

import javax.sql.DataSource;
import java.sql.SQLException;
import java.util.*;


/**
 * 解决5.x未配置分片的表无法指定dataSource
 * 思路:
 * 通过jpaEntityManager拿到所有实体类的表名, 不需要分片的表指定actual-node, 添加到shading的rule当中
 * jpaEntityManager拿所有实体依赖shardingSphereDataSource
 * shading的rule必须在shardingSphereDataSource创建之前指定，那上面的思路就是伪命题!
 * 但是强大的spring ioc中管理了dataSource, 既然不能在初始创建shardingSphereDataSource时拿到 所有的实体表名，那就另辟蹊径
 * <p>
 * 在shardingSphereDataSource创建好之后, 我在做一个桥接的bean,顺序在dataSource创建完成之后, 再获取到spring的bean工厂
 * 接下来就简单了, 之前的规则保留，再添加不需要分片表的规则，重新创建dataSource, 替换掉原ioc容器里的dataSource
 * 到此搞定！
 *
 * @author apelx
 * @since 2022-05-14
 */
@Configuration
@ComponentScan("org.apache.shardingsphere.spring.boot.converter")
@EnableConfigurationProperties(SpringBootPropertiesConfiguration.class)
@AutoConfigureAfter(value = ShardingSphereAutoConfiguration.class)
@RequiredArgsConstructor
public class SharingConfig implements BeanFactoryAware, EnvironmentAware {

    @Value("${sharding.default.datasource}")
    private String defaultDatasourceName;

    private String schemaName;

    private final SpringBootPropertiesConfiguration props;

    private final Map<String, DataSource> dataSourceMap = new LinkedHashMap<>();

    private DefaultListableBeanFactory beanFactory;


    @Bean
    @DependsOn(value = "entityBeanUtil")
    public ShardingPostProcess shardingPostProcess(final ObjectProvider<List<RuleConfiguration>> rules, final ObjectProvider<ModeConfiguration> modeConfig) throws SQLException {
        ShardingPostProcess shardingPostProcess = new ShardingPostProcess(defaultDatasourceName);

        // 后置处理添加自定义的规则，再重新构建dataSource
        Collection<RuleConfiguration> ruleConfigs = Optional.ofNullable(rules.getIfAvailable()).orElse(Collections.emptyList());
        ruleConfigs.forEach(shardingPostProcess::postProcess);
        DataSource dataSource = ShardingSphereDataSourceFactory.createDataSource(schemaName, modeConfig.getIfAvailable(), dataSourceMap, ruleConfigs, props.getProps());

        beanFactory.destroySingleton("shardingSphereDataSource");
        // 覆盖掉dataSource
        beanFactory.registerSingleton("shardingSphereDataSource", dataSource);

        return shardingPostProcess;
    }


    /**
     * Set the {@code Environment} that this component runs in.
     *
     * @param environment
     */

    @Override
    public void setEnvironment(Environment environment) {
        dataSourceMap.putAll(DataSourceMapSetter.getDataSourceMap(environment));
        schemaName = SchemaNameSetter.getSchemaName(environment);
    }

    @Override
    public void setBeanFactory(BeanFactory beanFactory) throws BeansException {
        this.beanFactory = (DefaultListableBeanFactory) beanFactory;
    }
}
